#include <opencv2/calib3d.hpp>
#include <opencv2/opencv.hpp>
#include <iostream> 
#include <chrono>
#include <random>
#include <vector>
#include <numeric>

// 生成矩形四个顶点的3D点
void generateRectangle3D(std::vector<cv::Point3f>& pts, float w, float h, float z) {
    pts.clear();
    pts.emplace_back(-w, -h, z);      // 左上
    pts.emplace_back(-w, h, z);      // 左下
    pts.emplace_back(w, h, z);      // 右下
    pts.emplace_back(w, -h, z);      // 右上
}

// 投影3D点到2D图像平面
void projectPoints(const std::vector<cv::Point3f>& objectPoints,
                   std::vector<cv::Point2f>& imagePoints,
                   const cv::Mat& rvec, const cv::Mat& tvec,
                   const cv::Mat& K, const cv::Mat& distCoeffs) {
    cv::projectPoints(objectPoints, rvec, tvec, K, distCoeffs, imagePoints);
}

// 给2D点加高斯噪声，模拟测量误差
// 参数：pts - 2D点，noise_sigma - 噪声标准差，gen - 随机数生成器
void addNoise(std::vector<cv::Point2f>& pts, float noise_sigma, std::mt19937& gen) {
    std::normal_distribution<float> dist(0, noise_sigma);
    for (auto& pt : pts) {
        pt.x += dist(gen); // x方向加噪声
        pt.y += dist(gen); // y方向加噪声
    }
}

// 评估位姿误差（旋转+平移）
// 参数：rvec_gt/tvec_gt - 真值，rvec/tvec - 估计值
double poseError(const cv::Mat& rvec_gt, const cv::Mat& tvec_gt, const cv::Mat& rvec, const cv::Mat& tvec) {
    // cv::Mat R_gt, R;
    // cv::Rodrigues(rvec_gt, R_gt); // 真值旋转矩阵
    // cv::Rodrigues(rvec, R);       // 估计旋转矩阵
    // double rot_err = cv::norm(R_gt - R, cv::NORM_L2); // 旋转误差
    double trans_err = cv::norm(tvec_gt - tvec);      // 平移误差
    return trans_err; // 误差
}

// 生成随机旋转向量和平移向量，模拟物体在空间中的随机姿态
// 参数：rvec/tvec - 输出旋转和平移，gen - 随机数生成器
void randomPose(cv::Mat& rvec, cv::Mat& tvec, std::mt19937& gen) {
    std::uniform_real_distribution<double> rot_distx(-0.46,0.86);   // 旋转范围
    std::uniform_real_distribution<double> rot_disty(-1.64,0.71);
    std::uniform_real_distribution<double> rot_distz(-0.10,0.09);
    std::uniform_real_distribution<double> trans_distx(-0.71,0.64);    // 平移范围
    std::uniform_real_distribution<double> trans_disty(0.0, 0.68);
    std::uniform_real_distribution<double> trans_distz(0.3, 2.36);
    rvec = (cv::Mat_<double>(3,1) << rot_distx(gen), rot_disty(gen), rot_distz(gen));
    tvec = (cv::Mat_<double>(3,1) << trans_distx(gen), trans_disty(gen), trans_distz(gen)); 
}

int main() {
    // 实验次数
    const int num_trials = 10000;
    float noise_sigma = 0.5f; // 噪声标准差
    float rect_w = 0.065f, rect_h = 0.0275f; // 矩形尺寸

    cv::Mat K = (cv::Mat_<double>(3,3) <<
            1720.623331817985,  0,                           641.1862135251633,
            0,                  1724.193394641095,           491.969263408209, 
            0,                  0,                           1);
    cv::Mat distCoeffs = (cv::Mat_<double>(5, 1) <<
    -0.0455753353769928, 0.3335555742001278, -0.00602518279942971, -0.007954908883944546, 0); // 畸变系数

    // 要测试的PnP方法
    std::vector<int> methods = {
        cv::SOLVEPNP_AP3P,
        cv::SOLVEPNP_IPPE,
        cv::SOLVEPNP_ITERATIVE
    };
    std::vector<std::string> method_names = {
        "AP3P", "IPPE","ITERATIVE"
    };

    // 记录每种方法的误差和耗时
    std::vector<std::vector<double>> errors(methods.size());
    std::vector<std::vector<double>> times(methods.size());
    std::vector<int> error_count_4cm(methods.size(), 0); // 记录误差大于10cm的次数

    std::mt19937 gen(1); // 随机种子

    // 主循环：多次实验
    for (int trial = 0; trial < num_trials; ++trial) {
        // 随机生成矩形在空间中的姿态
        std::vector<cv::Point3f> objectPoints;
        generateRectangle3D(objectPoints, rect_w, rect_h, 0.0f);

        cv::Mat rvec_gt, tvec_gt;
        randomPose(rvec_gt, tvec_gt, gen); // 随机生成真值位姿

        std::vector<cv::Point2f> imagePoints;
        projectPoints(objectPoints, imagePoints, rvec_gt, tvec_gt, K, distCoeffs); // 投影到2D
        addNoise(imagePoints, noise_sigma, gen); // 加噪声
        // 对所有二维点取整，模拟像素量化
        for (auto& pt : imagePoints) {
            pt.x = std::round(pt.x);
            pt.y = std::round(pt.y);
        }

        // 对每种PnP方法进行测试
        for (size_t i = 0; i < methods.size(); ++i) {
            std::vector<cv::Point3f> obj_pts = objectPoints;
            std::vector<cv::Point2f> img_pts = imagePoints;

            cv::Mat rvec, tvec;
            auto start = std::chrono::high_resolution_clock::now(); // 计时开始
            bool ok = cv::solvePnP(obj_pts, img_pts, K, distCoeffs, rvec, tvec, false, methods[i]); // 求解位姿
            auto end = std::chrono::high_resolution_clock::now();   // 计时结束
            double ms = std::chrono::duration<double, std::milli>(end-start).count(); // 计算耗时

            if (!ok || rvec.empty() || tvec.empty()) {
                continue; // 跳过异常样本
            }

            //打印出真实值与求解获得的值
            std::cout << "Method: " << method_names[i] << std::endl;
            std::cout << "True rvec: " << rvec_gt.t() << std::endl;
            std::cout << "Est  rvec: " << rvec.t() << std::endl;
            std::cout << "True tvec: " << tvec_gt.t() << std::endl;
            std::cout << "Est  tvec: " << tvec.t() << std::endl;

            double err = ok ? poseError(rvec_gt, tvec_gt, rvec, tvec) : -1; // 计算误差
            if (ok) {
                errors[i].push_back(err); // 记录误差
                times[i].push_back(ms);   // 记录耗时
                if (err > 0.04) error_count_4cm[i]++; // 误差大于4cm计数
            }
        }
    }

    // 输出每种方法的均值和方差（均值仅统计误差小于4cm的样本）
    for (size_t i = 0; i < methods.size(); ++i) {
        // 仅统计误差小于4cm的样本
        std::vector<double> filtered_errors;
        for (auto v : errors[i]) {
            if (v < 0.04 && v >= 0) filtered_errors.push_back(v);
        }
        double mean_err = filtered_errors.empty() ? -1 :
            std::accumulate(filtered_errors.begin(), filtered_errors.end(), 0.0) / filtered_errors.size();
        double mean_time = times[i].empty() ? -1 :
            std::accumulate(times[i].begin(), times[i].end(), 0.0) / times[i].size();   // 平均耗时

        double std_err = 0, std_time = 0;
        if (!filtered_errors.empty()) {
            for (auto v : filtered_errors) std_err += (v - mean_err) * (v - mean_err);
            std_err = sqrt(std_err / filtered_errors.size()); // 误差标准差
        }
        if (!times[i].empty()) {
            for (auto v : times[i]) std_time += (v - mean_time) * (v - mean_time);
            std_time = sqrt(std_time / times[i].size()); // 耗时标准差
        }

        // 打印结果
        std::cout << method_names[i] << ":\tmean_error(<4cm)=" << mean_err
                  << "\tstd_error(<4cm)=" << std_err
                  << "\tmean_time=" << mean_time
                  << " ms\tstd_time=" << std_time << " ms\t"
                  << "samples=" << errors[i].size()
                  << "\terror>4cm=" << error_count_4cm[i] << std::endl;
    }

    return 0;
}